#!/usr/bin/env python3
"""
翻译功能接口 (Translation API Interface)
"""

import time
import re
from typing import Optional, Dict, List
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field
from vllm import SamplingParams

# 全局模型实例（从main.py导入）
model = None

# 支持的语言代码及中文含义
SUPPORTED_LANGUAGES = {
    "zh": "中文",
    "en": "英语", 
    "ja": "日语",
    "ko": "韩语",
    "fr": "法语",
    "de": "德语",
    "es": "西班牙语",
    "it": "意大利语",
    "pt": "葡萄牙语",
    "ru": "俄语",
    "ar": "阿拉伯语",
    "hi": "印地语",
    "th": "泰语",
    "vi": "越南语",
    "id": "印尼语",
    "ms": "马来语",
    "tr": "土耳其语",
    "pl": "波兰语",
    "nl": "荷兰语",
    "sv": "瑞典语",
    "da": "丹麦语",
    "no": "挪威语",
    "fi": "芬兰语",
    "cs": "捷克语",
    "hu": "匈牙利语",
    "ro": "罗马尼亚语",
    "bg": "保加利亚语",
    "hr": "克罗地亚语",
    "sk": "斯洛伐克语",
    "sl": "斯洛文尼亚语",
    "et": "爱沙尼亚语",
    "lv": "拉脱维亚语",
    "lt": "立陶宛语",
    "uk": "乌克兰语",
    "be": "白俄罗斯语",
    "mk": "马其顿语",
    "sq": "阿尔巴尼亚语",
    "sr": "塞尔维亚语",
    "bs": "波斯尼亚语",
    "me": "黑山语",
    "is": "冰岛语",
    "mt": "马耳他语",
    "cy": "威尔士语",
    "ga": "爱尔兰语",
    "gd": "苏格兰盖尔语",
    "eu": "巴斯克语",
    "ca": "加泰罗尼亚语",
    "gl": "加利西亚语",
    "ast": "阿斯图里亚斯语",
    "oc": "奥克语",
    "br": "布列塔尼语",
    "co": "科西嘉语",
    "sc": "撒丁语",
    "rm": "罗曼什语",
    "lb": "卢森堡语"
}

# 语言检测模式
LANGUAGE_DETECTION_PATTERNS = {
    "zh": r'[\u4e00-\u9fff]',
    "ja": r'[\u3040-\u309f\u30a0-\u30ff]',
    "ko": r'[\uac00-\ud7af]',
    "ar": r'[\u0600-\u06ff]',
    "th": r'[\u0e00-\u0e7f]',
    "hi": r'[\u0900-\u097f]',
    "ru": r'[\u0400-\u04ff]',
}

class TranslateRequest(BaseModel):
    """翻译请求模型 (Translation request model)"""
    text: str = Field(..., description="待翻译内容 (Text to translate)", min_length=1, max_length=5000)
    target_language: str = Field("zh", description="目标语言代码 (Target language code)")
    source_language: Optional[str] = Field(None, description="源语言代码，为空时自动检测 (Source language code, auto-detect if empty)")

class TranslateResponse(BaseModel):
    """翻译响应模型 (Translation response model)"""
    original_text: str = Field(..., description="翻译前内容 (Original text)")
    translated_text: str = Field(..., description="翻译后内容 (Translated text)")
    source_language_code: str = Field(..., description="源语言代码 (Source language code)")
    source_language_name: str = Field(..., description="源语言含义 (Source language name)")
    target_language_code: str = Field(..., description="目标语言代码 (Target language code)")
    target_language_name: str = Field(..., description="目标语言含义 (Target language name)")
    translation_time: float = Field(..., description="翻译耗时(秒) (Translation time in seconds)")

class LanguageInfo(BaseModel):
    """语言信息模型 (Language info model)"""
    code: str = Field(..., description="语言代码 (Language code)")
    name: str = Field(..., description="语言中文名称 (Chinese name of language)")

class SupportedLanguagesResponse(BaseModel):
    """支持语言列表响应模型 (Supported languages response model)"""
    languages: List[LanguageInfo] = Field(..., description="支持的语言列表 (List of supported languages)")
    total_count: int = Field(..., description="支持语言总数 (Total count of supported languages)")

def detect_language(text: str) -> str:
    """
    检测文本语言 (Detect text language)
    
    Args:
        text: 待检测文本 (Text to detect)
        
    Returns:
        str: 检测到的语言代码 (Detected language code)
    """
    # 简单的语言检测逻辑
    for lang_code, pattern in LANGUAGE_DETECTION_PATTERNS.items():
        if re.search(pattern, text):
            return lang_code
    
    # 默认检测为英语
    return "en"

def create_translation_prompt(text: str, source_lang: str, target_lang: str) -> str:
    """
    创建翻译提示词 (Create translation prompt)
    
    Args:
        text: 待翻译文本 (Text to translate)
        source_lang: 源语言代码 (Source language code)
        target_lang: 目标语言代码 (Target language code)
        
    Returns:
        str: 翻译提示词 (Translation prompt)
    """
    source_name = SUPPORTED_LANGUAGES.get(source_lang, "未知语言")
    target_name = SUPPORTED_LANGUAGES.get(target_lang, "未知语言")
    
    prompt = f"""请将以下{source_name}文本翻译成{target_name}，只返回翻译结果，不要添加任何解释或其他内容：

原文：{text}

翻译："""
    
    return prompt

def extract_translation_result(generated_text: str, original_text: str) -> str:
    """
    从生成的文本中提取翻译结果 (Extract translation result from generated text)
    
    Args:
        generated_text: 模型生成的文本 (Generated text from model)
        original_text: 原始文本 (Original text)
        
    Returns:
        str: 提取的翻译结果 (Extracted translation result)
    """
    # 移除可能的前缀和后缀
    result = generated_text.strip()
    
    # 移除常见的前缀
    prefixes_to_remove = ["翻译：", "Translation:", "译文：", "结果：", "Result:"]
    for prefix in prefixes_to_remove:
        if result.startswith(prefix):
            result = result[len(prefix):].strip()
            break
    
    # 如果结果为空或与原文相同，返回原文
    if not result or result == original_text:
        return original_text
    
    return result

# 创建路由器
router = APIRouter(prefix="/translate", tags=["翻译 (Translation)"])

@router.post("/", response_model=TranslateResponse)
async def translate_text(request: TranslateRequest):
    """
    翻译文本 (Translate text)
    
    Args:
        request: 翻译请求 (Translation request)
        
    Returns:
        TranslateResponse: 翻译响应 (Translation response)
    """
    global model
    
    if model is None:
        raise HTTPException(status_code=503, detail="模型未加载 (Model not loaded)")
    
    # 验证目标语言代码
    if request.target_language not in SUPPORTED_LANGUAGES:
        raise HTTPException(
            status_code=400, 
            detail=f"不支持的目标语言代码: {request.target_language} (Unsupported target language code)"
        )
    
    start_time = time.time()
    
    try:
        # 检测或验证源语言
        if request.source_language:
            if request.source_language not in SUPPORTED_LANGUAGES:
                raise HTTPException(
                    status_code=400,
                    detail=f"不支持的源语言代码: {request.source_language} (Unsupported source language code)"
                )
            source_lang = request.source_language
        else:
            source_lang = detect_language(request.text)
        
        # 如果源语言和目标语言相同，直接返回原文
        if source_lang == request.target_language:
            translation_time = time.time() - start_time
            return TranslateResponse(
                original_text=request.text,
                translated_text=request.text,
                source_language_code=source_lang,
                source_language_name=SUPPORTED_LANGUAGES[source_lang],
                target_language_code=request.target_language,
                target_language_name=SUPPORTED_LANGUAGES[request.target_language],
                translation_time=translation_time
            )
        
        # 创建翻译提示词
        prompt = create_translation_prompt(request.text, source_lang, request.target_language)
        
        # 设置采样参数
        sampling_params = SamplingParams(
            temperature=0.1,  # 较低的温度以获得更一致的翻译
            max_tokens=min(len(request.text) * 3, 1024),  # 根据输入长度动态调整
            skip_special_tokens=True,
            top_p=0.9,
            repetition_penalty=1.05
        )
        
        # 生成翻译
        results = model.generate([prompt], sampling_params)
        generated_text = results[0].outputs[0].text.strip()
        
        # 提取翻译结果
        translated_text = extract_translation_result(generated_text, request.text)
        
        translation_time = time.time() - start_time
        
        return TranslateResponse(
            original_text=request.text,
            translated_text=translated_text,
            source_language_code=source_lang,
            source_language_name=SUPPORTED_LANGUAGES[source_lang],
            target_language_code=request.target_language,
            target_language_name=SUPPORTED_LANGUAGES[request.target_language],
            translation_time=translation_time
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"翻译失败 (Translation failed): {str(e)}")

@router.get("/languages", response_model=SupportedLanguagesResponse)
async def get_supported_languages():
    """
    获取所有支持的语言代码及中文含义 (Get all supported language codes and Chinese meanings)
    
    Returns:
        SupportedLanguagesResponse: 支持的语言列表 (List of supported languages)
    """
    languages = [
        LanguageInfo(code=code, name=name) 
        for code, name in SUPPORTED_LANGUAGES.items()
    ]
    
    # 按语言代码排序
    languages.sort(key=lambda x: x.code)
    
    return SupportedLanguagesResponse(
        languages=languages,
        total_count=len(languages)
    )

@router.get("/detect")
async def detect_text_language(text: str):
    """
    检测文本语言 (Detect text language)
    
    Args:
        text: 待检测文本 (Text to detect)
        
    Returns:
        dict: 检测结果 (Detection result)
    """
    if not text or len(text.strip()) == 0:
        raise HTTPException(status_code=400, detail="文本不能为空 (Text cannot be empty)")
    
    detected_lang = detect_language(text)
    
    return {
        "text": text,
        "detected_language_code": detected_lang,
        "detected_language_name": SUPPORTED_LANGUAGES.get(detected_lang, "未知语言"),
        "confidence": "简单检测" if detected_lang != "en" else "默认检测"
    }

def set_model(model_instance):
    """
    设置全局模型实例 (Set global model instance)
    
    Args:
        model_instance: 模型实例 (Model instance)
    """
    global model
    model = model_instance
